#!/usr/bin/env python
import os
import argparse
import numpy as np
from astropy.io import fits
from scipy.ndimage import binary_dilation, binary_erosion
from scipy.ndimage import median_filter
from scipy.optimize import curve_fit
from matplotlib import pyplot as plt
from mpl_toolkits.axes_grid1 import make_axes_locatable


def gauss(x, *p):
    a, mu, sigma = p
    return a * np.exp(-(x - mu) ** 2 / (2. * sigma ** 2))


def xy2vec(xx, yy):
    vec = np.array([xx, yy])
    vec = np.swapaxes(vec, 0, -1)
    return vec


def dist2line(p1, p2, vec):
    v0 = p2 - p1
    return np.cross(v0, vec - p1) / np.linalg.norm(v0)


def get_mtf(data, pixel_size, oversample=5, freq=50.0, ffig=None):
    ny, nx = data.shape
    xx, yy = np.meshgrid(np.arange(nx), np.arange(ny))

    vmax = np.percentile(data, 95)
    vmin = np.percentile(data, 5)
    thresh = (vmin + vmax) / 2
    index = data > thresh
    edge1 = binary_dilation(index) * (~index)
    edge2 = binary_erosion(index) * (~index)
    edge = np.logical_or(edge1, edge2)

    edge_x, edge_y = xx[edge], yy[edge]
    pfit = np.polyfit(edge_x, edge_y, 1)

    x1 = nx // 2
    y1 = x1 * pfit[0] + pfit[1]
    p1 = np.array([x1, y1])
    x2 = 0
    y2 = pfit[1]
    p2 = np.array([x2, y2])

    vec = xy2vec(xx.ravel(), yy.ravel())
    d = dist2line(p1, p2, vec)
    v = data.ravel()
    if v[d > 0].mean() < v[d < 0].mean():
        d = d * -1

    dmin, dmax = min(d), max(d)
    dmin_outer = int(np.floor(dmin))
    dmax_outer = int(np.floor(dmax)) + 1
    n = (dmax_outer - dmin_outer) * oversample
    xlims = np.arange(n + 1) / oversample + dmin_outer
    x = (np.arange(n) + 0.5) / oversample + dmin_outer

    y = np.full(n, np.nan)
    for i in range(n):
        index = (d > xlims[i]) * (d <= xlims[i + 1])
        if np.sum(index) > 0:
            y[i] = np.mean(v[index])

    # fill missing data
    index = np.isfinite(y)
    y = np.interp(x, x[index], y[index], left=np.nan, right=np.nan)
    index = np.isfinite(y)
    x = x[index]
    y = y[index]

    ydiff = y[1:] - y[:-1]
    xdiff = (x[1:] + x[:-1]) / 2
    p0 = [np.max(y), 0., 0.5]
    coeff, var_matrix = curve_fit(gauss, xdiff, ydiff, p0=p0)

    # calculate mtf
    amp = np.absolute(np.fft.fft(ydiff))
    ff = np.fft.fftfreq(len(amp), d=pixel_size / 1000 / oversample)
    amp = amp / amp[0]
    fmax = min(np.max(ff), freq * 2)
    index = (ff >= 0) * (ff < fmax)
    amp = amp[index]
    ff = ff[index]
    mtf = np.interp(freq, ff, amp)

    if ffig is not None:
        plt.figure(figsize=(11, 10))
        plt.subplot(221)
        im = plt.imshow(data)
        ax = plt.gca()
        divider = make_axes_locatable(ax)
        cax = divider.append_axes('right', size='5%', pad=0.05)
        plt.colorbar(im, cax=cax)

        plt.subplot(222)
        plt.scatter(x, y, color='g', s=5)

        plt.subplot(223)
        plt.scatter(xdiff, ydiff, color='g', s=5)
        l, = plt.plot(xdiff, gauss(xdiff, *coeff), color='orange', ls='--')
        plt.legend([l, ], ['{:.2f}{}{:.2f}'.format(coeff[1], chr(177), coeff[2])])

        plt.subplot(224)
        plt.plot(ff, amp, color='g', marker='o')
        ax = plt.gca()
        ylim = ax.get_ylim()
        l, = plt.plot([freq, freq], ylim, color='orange', ls='--')
        plt.ylim(ylim)
        plt.legend([l, ], ['MTF@{:.1f}lp/mm = {:.3f}'.format(freq, mtf)])

        plt.tight_layout()
        plt.savefig(ffig)
        plt.close()

    return mtf, coeff[2]


def measure_mtf(fdata, fbias=None, xcent=None, ycent=None, xsize=50, ysize=50,
                smooth=2, pixelsize=10.0, oversample=5, freq=20.0, ffig=None,num_out_flag=False):
    with fits.open(fdata) as img:
        image = img[0].data.astype(np.float32)
        if fbias is not None:
            image = image - fits.getdata(fbias).astype(np.float32)
        if len(image.shape) == 3:
            image = image.mean(axis=0)  # merge color channels
        if smooth > 0:
            image = median_filter(image, smooth * 2 + 1)

        # cutout
        ny, nx = image.shape
        if xcent is None:
            xmin, xmax = 0, nx
        else:
            xmin, xmax = xcent - xsize // 2, xcent + xsize // 2
        if ycent is None:
            ymin, ymax = 0, ny
        else:
            ymin, ymax = ycent - ysize // 2, ycent + ysize // 2
        image = image[ymin:ymax, xmin:xmax]

        mtf, sigma = get_mtf(image, pixelsize, oversample=oversample, freq=freq, ffig=ffig)
        
        if num_out_flag == False:
            return ['{:.2f}'.format(freq), '{:.4f}'.format(mtf), '{:.4f}'.format(sigma)]
        else:
            return mtf, sigma


def _get_seqnum(name):
    with fits.open(name) as img:
        if 'SEQNUM' in img[0].header:
            return img[0].header['SEQNUM']
        else:
            return 'name' + name


def _write_log(flog, fname, result):
    if os.path.exists(flog) is False:
        with open(flog, 'w') as f:
            f.write('# name, frequency, mtf\n')
    with open(flog, 'a') as f:
        f.write(fname + ', ' + ', '.join(result) + '\n')


def _run(flist, fbias, xcent, ycent, xsize, ysize,
         smooth, pixelsize, oversample, freq, flog, debug):
    # sort by seqnum
    flist.sort(key=_get_seqnum)
    # loop
    for f in flist:
        if debug:
            ffig = f.replace('.fits', '.png')
        else:
            ffig = None
        res = measure_mtf(f, fbias=fbias,
                          xcent=xcent, ycent=ycent, xsize=xsize, ysize=ysize,
                          smooth=smooth, pixelsize=pixelsize, oversample=oversample,
                          freq=freq, ffig=ffig)
        print(res)
        res += ['{:.2f}'.format(xcent),'{:.2f}'.format(ycent),'{:.2f}'.format(xsize),'{:.2f}'.format(ysize)]
        if flog is not None:
            _write_log(flog, f, res)
    

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('fname', type=str, nargs='*')
    parser.add_argument('--bias', default=None, type=str,
                        help='bias file name')
    parser.add_argument('--xcent', '-x', default=None, type=int,
                        help='x position of cutout center. If not set, use the full image')
    parser.add_argument('--ycent', '-y', default=None, type=int,
                        help='y position of cutout center. If not set, use the full image')
    parser.add_argument('--xsize', default=50, type=int,
                        help='x size of cutout, default is 50')
    parser.add_argument('--ysize', default=50, type=int,
                        help='y size of cutout, default is 50')
    parser.add_argument('--smooth', '-m', default=2, type=int,
                        help='smooth width = SMOOTH * 2 + 1, default is 2')
    parser.add_argument('--pixelsize', '-p', default=10, type=float,
                        help='pixel size in um, default is 10')
    parser.add_argument('--oversample', '-s', default=5, type=int,
                        help='oversample value for projection profile, default is 5')
    parser.add_argument('--freq', '-f', default=20.0, type=float,
                        help='frequency for MTF evaluation in lp/mm, default is 20')
    parser.add_argument('--log', default=None, type=str,
                        help='output result table file')
    parser.add_argument('--debug', action='store_true',
                        help='if set, make debug figure')
    pars = parser.parse_args()
    _run(pars.fname, pars.bias, pars.xcent, pars.ycent, pars.xsize, pars.ysize,
         pars.smooth, pars.pixelsize, pars.oversample, pars.freq, pars.log, pars.debug)
